package com.api.apigateway.filter;

import cn.hutool.json.JSONUtil;
import com.api.common.common.ErrorCode;
import com.api.common.common.ResultUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Map;

@Slf4j
@Component
public class GatewayRequestRateLimiterGatewayFilterFactory extends RequestRateLimiterGatewayFilterFactory {

    private final RateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;


    private static final String BLACKLIST_KEY_PREFIX = "blacklist:";
    public GatewayRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver) {
        super(defaultRateLimiter, defaultKeyResolver);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;

    }

    @Override
    public GatewayFilter apply(Config config) {
        log.info("-----------GatewayRequestRateLimiterGatewayFilterFactory------------");
        KeyResolver resolver = getOrDefault(config.getKeyResolver(), defaultKeyResolver);
        RateLimiter<Object> limiter = getOrDefault(config.getRateLimiter(), defaultRateLimiter);
        //resolver.resolve(exchange) 获取限流的 Key
        return (exchange, chain) -> resolver.resolve(exchange).flatMap(key -> {
            String routeId = config.getRouteId();
            if (routeId == null) {
                //动态获取路由 ID
                Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                routeId = route.getId();
            }
            String finalRouteId = routeId;
            //调用 limiter.isAllowed(routeId, key) 来判断该用户是否被允许继续请求
            //如果限流被触发（response.isAllowed() 为 false），你自定义了返回的响应，状态码为 500 Internal Server Error，并返回自定义的 JSON 消息。
            return limiter.isAllowed(routeId, key).flatMap(response -> {
                //todo 可以添加对黑名单的校验
//                return redisTemplate.hasKey(BLACKLIST_KEY_PREFIX + key).flatMap(isBlacklisted -> {
//                    if (isBlacklisted) {
//                        return rejectRequest(exchange.getResponse(), "您已被加入黑名单");
//                    }
//                    // 原有限流逻辑
//                });
                for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                    exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                }
                if (response.isAllowed()) {
                    return chain.filter(exchange);
                }
                log.warn("已限流: {}", finalRouteId);
                ServerHttpResponse httpResponse = exchange.getResponse();
                return rejectRequest(httpResponse);

            });
        });
    }

    private static Mono<Void> rejectRequest(ServerHttpResponse httpResponse) {
        //修改code为429
        httpResponse.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        if (!httpResponse.getHeaders().containsKey("Content-Type")) {
            httpResponse.getHeaders().add("Content-Type", "application/json");
        }
        byte[] bytes = JSONUtil.toJsonStr(ResultUtils.error(ErrorCode.INVOKE_INTERFACE_ERROR, "请求次数过多，请稍后再试。")).getBytes();
        httpResponse.getHeaders().setContentLength(bytes.length);
        //此处无法触发全局异常处理，手动返回
        DataBuffer buffer = httpResponse.bufferFactory().wrap(bytes);
        return httpResponse.writeWith(Mono.just(buffer));
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }
}

